import os
import json
from datetime import datetime
from io import StringIO

import pandas as pd
from flask import Flask, request, jsonify, render_template, send_from_directory, Response, abort
import joblib

MODELS_DIR = "models"
ASSETS_DIR = "assets"

app = Flask(__name__, template_folder="templates", static_folder=None)

def load_artifacts():
    """
    加载模型和元数据
    """
    model_path = os.path.join(MODELS_DIR, "california_housing_pipeline.joblib")
    features_path = os.path.join(MODELS_DIR, "feature_names.json")
    metrics_path = os.path.join(MODELS_DIR, "metrics.json")

    if not os.path.exists(model_path):
        raise FileNotFoundError("Model not found. Please run: python train.py")
    
    model = joblib.load(model_path)
    with open(features_path, "r", encoding="utf-8") as f:
        feature_names = json.load(f)

    metrics = {}
    if os.path.exists(metrics_path):
        with open(metrics_path, "r", encoding="utf-8") as f:
            metrics = json.load(f)

    return model, feature_names, metrics

model, FEATURE_NAMES, METRICS = load_artifacts()

@app.get("/")
def index():
    return render_template("index.html", feature_names=FEATURE_NAMES, metrics=METRICS)

@app.get("/health")
def health():
    return jsonify({"status": "ok", "time": datetime.now().isoformat() + "Z"})

@app.get("/metrics")
def metrics():
    return jsonify(METRICS)

@app.get("/plots/<path:filename>")
def plots(filename: str):
    # 提供训练时保存的图像
    # safe = {"pred_vs_true.png", "residuals.png", "feature_importance.png"}
    safe = {"pred_vs_true.png", "residuals.png"}
    if filename not in safe:
        abort(404)
    return send_from_directory(ASSETS_DIR, filename)

def to_dataframe_from_json(payload):
    # 支持两种 JSON 结构：单对象 或 对象数组
    if isinstance(payload, dict):
        df = pd.DataFrame([payload], columns=FEATURE_NAMES)
    elif isinstance(payload, list):
        df = pd.DataFrame(payload, columns=FEATURE_NAMES)
    else:
        raise ValueError("JSON payload must be an object or a list of objects.")
    # 转为浮点
    for col in FEATURE_NAMES:
        df[col] = pd.to_numeric(df[col], errors="coerce")
    if df[FEATURE_NAMES].isnull().any().any():
        missing_rows = df[FEATURE_NAMES].isnull().any(axis=1).sum()
        raise ValueError(f"Found non-numeric or missing values in {missing_rows} row(s).")
    return df

@app.post("/predict")
def predict_json():
    try:
        payload = request.get_json(force=True, silent=False)
        df = to_dataframe_from_json(payload)
        preds = model.predict(df)
        return jsonify({
            "predictions": preds.tolist(),
            "count": len(preds),
            "feature_order": FEATURE_NAMES
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 400

@app.post("/predict-csv")
def predict_csv():
    """
    上传 CSV，列名必须包含 FEATURE_NAMES（可以有多余列），返回附带 prediction 列的 CSV
    """
    if "file" not in request.files:
        return jsonify({"error": "Missing file field 'file'."}), 400
    file = request.files["file"]
    if file.filename == "":
        return jsonify({"error": "Empty filename."}), 400
    try:
        content = file.read().decode("utf-8", errors="ignore")
        df_in = pd.read_csv(StringIO(content))
        missing = [c for c in FEATURE_NAMES if c not in df_in.columns]
        if missing:
            return jsonify({"error": f"Missing required columns: {missing}"}), 400
        df = df_in[FEATURE_NAMES].copy()
        for col in FEATURE_NAMES:
            df[col] = pd.to_numeric(df[col], errors="coerce")
        if df.isnull().any().any():
            return jsonify({"error": "Found non-numeric or missing values in required columns."}), 400
        preds = model.predict(df)
        out = df_in.copy()
        out["prediction"] = preds
        csv_buf = StringIO()
        out.to_csv(csv_buf, index=False)
        csv_buf.seek(0)
        return Response(
            csv_buf.getvalue(),
            mimetype="text/csv",
            headers={"Content-Disposition": "attachment; filename=predictions.csv"}
        )
    except Exception as e:
        return jsonify({"error": str(e)}), 400

if __name__ == "__main__":
    # 直接运行：python app.py
    app.run(host="0.0.0.0", port=5000, debug=True)